use std::io::Write;
use std::path::Path;
use std::fs::{self, File};



// ------------ Parser -------------------

#[derive(Debug)]
enum Segment {
    Argument, Local, That, This,
    Pointer, Temp,
    Constant,
    Static,
}

#[derive(Debug)]
enum Code {
    Add, Sub, Neg, 
    Eq,  Gt,  Lt,
    And, Or,  Not,
    Push(Segment, u16),
    Pop(Segment, u16),
    Label(String),
    Goto(String),
    IfGoto(String),
    Function(String, u16),
    Call(String, u16),
    Return,
}

use Code::*;
use Segment::*;
struct Parser { content: String }

impl Parser {

    fn new(file: &str) -> Self {
        let content = fs::read_to_string(file).unwrap();
        Parser { content }
    }

    fn parse(self) -> Vec<Code> {
        let code = self.filter_comment_line();
        code.into_iter()
            .map(|line| Self::parse_single_line(line))
            .collect()
    }

    fn filter_comment_line(self) -> Vec<String> {
        let codes = self.content.split("\n")
            .filter(|x| !(x.starts_with("//") || x.trim() == ""))
            .map(|x| x.to_string())
            .collect();
        return codes;
    }

    fn parse_single_line(line: String) -> Code {
        let strlist: Vec<&str> = line.as_str().trim().split(|c| c == ' ' || c == '\t').collect();
        match strlist[0] {
            "add" => Add, "sub" => Sub, "neg" => Neg,
            "eq"  => Eq,  "gt"  => Gt,  "lt"  => Lt,
            "and" => And, "or"  => Or,  "not" => Not,
            "push" | "pop" => {
                let segm = match strlist[1] {
                    "local" => Local,
                    "argument" => Argument,
                    "that" => That,
                    "this" => This,
                    "pointer" => Pointer,
                    "temp" => Temp,
                    "constant" => Constant,
                    "static" => Static,
                    _ => panic!("Invalid Segment!"),
                };
                let index: u16 = strlist[2].trim().parse().unwrap();
                if strlist[0] == "push" {
                    return Push(segm, index);
                }
                return Pop(segm, index);
            }
            "label" => Label(strlist[1].to_string()),
            "goto" => Goto(strlist[1].to_string()),
            "if-goto" => IfGoto(strlist[1].to_string()),
            "function" => Function(strlist[1].to_string(), strlist[2].trim().parse().unwrap()),
            "call" => Call(strlist[1].to_string(), strlist[2].trim().parse().unwrap()),
            "return" => Return,
            _ => panic!("Invalid code"),
        }
    }
}

// ------------ Code Writer -------------------

fn handle_comp(comp: &str, idx: usize) -> String {
    let res = "@SP    \n".to_owned()
            + "AM=M-1 \n"
            + "D=M    \n"
            + "A=A-1  \n"
            + "D=M-D  \n";
    // jump here
    let true_i = format!("TRUE{}", idx);
    let continue_i = format!("CONTINUE{}", idx);
    let jump = format!( 
       "@{}
        D;{}
        D=0
        @{}
        0;JMP
    ({})
        D=-1
    ({})
        @SP
        AM=M-1
        M=D
        @SP
        M=M+1
    ", true_i, comp, continue_i, true_i, continue_i);
    return res + &jump;
}

/// main generate code function
/// # Args
/// * `code`: the Code
/// * `fname`: label use to handle static variable
/// * `idx`: index use to handle comp jump
fn generate_single_code(code: Code, fname: &str, idx: usize) -> String {
    match code {
        Push(segm, index) => {
            let offset = format!("@{}   \n", index)
                                + "D=A   \n";
            let set_data = match segm {
                Constant => offset,
                Argument => offset + "@ARG   \n  A=D+M  \n D=M \n",
                Local    => offset + "@LCL   \n  A=D+M  \n D=M \n",
                This     => offset + "@THIS  \n  A=D+M  \n D=M \n",
                That     => offset + "@THAT  \n  A=D+M  \n D=M \n",
                Pointer  => offset + "@3     \n  A=D+A  \n D=M \n",
                Temp     => offset + "@5     \n  A=D+A  \n D=M \n",
                Static   => format!( "@{}.{} \n", fname, index) + "D=M \n",
            };
            let in_stack = "@SP   \n\
                            A=M   \n\
                            M=D   \n\
                            @SP   \n\
                            M=M+1 \n";
            return set_data + in_stack;
        }
        Pop(segm, index) => {
            let offset = format!("@{}   \n", index)
                                + "D=A   \n";
            let set_data = "@R13    \n\
                            M=D     \n\
                            @SP     \n\
                            AM=M-1  \n\
                            D=M     \n\
                            @R13    \n\
                            A=M     \n\
                            M=D     \n";
            match segm {
                Argument => offset + "@ARG   \n  D=D+M   \n" + set_data,
                Local    => offset + "@LCL   \n  D=D+M   \n" + set_data,
                This     => offset + "@THIS  \n  D=D+M   \n" + set_data,
                That     => offset + "@THAT  \n  D=D+M   \n" + set_data,
                Pointer  => offset + "@3     \n  D=D+A   \n" + set_data,
                Temp     => offset + "@5     \n  D=D+A   \n" + set_data,
                Static   => format!( "@SP    \n\
                                        AM=M-1 \n\
                                        D=M    \n\
                                        @{}.{} \n\
                                        M=D    \n", fname, index),
                Constant => panic!("Could not pop on Constant Segment"),
            }
        }
        Add | Sub | Or | And => {
            let c1 = "@SP     \n".to_owned() 
                    + "AM=M-1  \n"
                    + "D=M     \n"
                    + "A=A-1   \n";
            let c2 = match code {
                Add => "M=D+M \n",
                Sub => "M=M-D \n",
                Or  => "M=D|M \n",
                And => "M=D&M \n",
                _ => unreachable!(),
            };
            return c1 + c2;
        }
        Neg | Not => {
            let c1 = "@SP   \n".to_owned() 
                    + "AM=M-1\n";
            let c2 = match code {
                Neg => "M=-M\n",
                Not => "M=!M\n",
                _ => unreachable!(),
            };
            let c3 = "@SP   \n";
            let c4 = "M=M+1 \n";
            return c1 + c2 + c3 + c4;
        }
        Eq => handle_comp("JEQ", idx),
        Gt => handle_comp("JGT", idx),
        Lt => handle_comp("JLT", idx),
        Label(label) => format!("({}${})", fname, label),
        Goto(label) => format!("@{}${}  \n  0;JMP   \n", fname, label),
        IfGoto(label) => format!(
            "@SP       \n\
            AM=M-1    \n\
            D=M       \n\
            @{}${}    \n\
            D;JNE     \n", fname, label),
        Function(name, nlocals) => {
            let mut fun = format!("({})    \n", name);
            for _i in 0..nlocals {
                let push0 = "@SP    \n\
                                A=M    \n\
                                M=0    \n\
                                @SP    \n\
                                M=M+1  \n";
                fun += push0;
            }
            return fun;
        },
        Call(name, argc) => {
            let ret_addr = format!("RET_{}_{}", name, idx);
            let push = |sym| {
                format!("@{}    \n\
                            D=M    \n\
                            @SP    \n\
                            A=M    \n\
                            M=D    \n\
                            @SP    \n\
                            M=M+1  \n", sym)
            };
            let push_ret = format!(
                        "@{}    \n\
                        D=A    \n\
                        @SP    \n\
                        A=M    \n\
                        M=D    \n\
                        @SP    \n\
                        M=M+1  \n", ret_addr);
            let push4 = format!("{}{}{}{}", 
                push("LCL"),  push("ARG"),
                push("THIS"), push("THAT")
            );
            let offset = argc + 5;
            let init = format!(
                        "@{}     \n\
                        D=A     \n\
                        @SP     \n\
                        D=M-D   \n\
                        @ARG    \n\
                        M=D     \n\
                        @SP     \n\
                        D=M     \n\
                        @LCL    \n\
                        M=D     \n\
                        @{}     \n\
                        0;JMP   \n\
                        ({})    \n", offset, name, ret_addr);
            return push_ret + &push4 + &init;
        }
        Return => {
            let save = "@LCL    \n\
                        D=M     \n\
                        @R13    \n\
                        M=D     \n".to_owned();     // Frame in R13
            let reta = "@5      \n\
                        D=A     \n\
                        @R13    \n\
                        A=M-D   \n\
                        D=M     \n\
                        @R14    \n\
                        M=D     \n";    // ret addr in R14
            let retv = "@SP     \n\
                        AM=M-1  \n\
                        D=M     \n\
                        @ARG    \n\
                        A=M     \n\
                        M=D     \n\
                        @ARG    \n\
                        D=M+1   \n\
                        @SP     \n\
                        M=D     \n";    // set return val on stack pop
            let that = "@R13    \n\
                        A=M-1   \n\
                        D=M     \n\
                        @THAT   \n\
                        M=D     \n";
            let this = "@R13    \n\
                        D=M     \n\
                        @2      \n\
                        A=D-A   \n\
                        D=M     \n\
                        @THIS   \n\
                        M=D     \n";
            let arg  = "@R13    \n\
                        D=M     \n\
                        @3      \n\
                        A=D-A   \n\
                        D=M     \n\
                        @ARG    \n\
                        M=D     \n";
            let lcl  = "@R13    \n\
                        D=M     \n\
                        @4      \n\
                        A=D-A   \n\
                        D=M     \n\
                        @LCL    \n\
                        M=D     \n";
            let goto = "@R14    \n\
                        A=M     \n\
                        0;JMP   \n";
            return save + reta + retv + that + this + arg + lcl + goto;
        }
    }
}

/// generate SINGLE asm file from possibly many files 
/// # Args
/// - `codes`: codes from many files
/// - `fnames`: the files name stem
/// - `has_boot`: whether to include bootstrap code
fn generate_codes(codes: Vec<Vec<Code>>,  fnames: Vec<&str>, has_boot: bool) -> Vec<String> {
    let mut idx = 0;
    let mut gencodes = vec![];
    if has_boot {
        let set_sp = "@256        \n\
                      D=A         \n\
                      @SP         \n\
                      M=D         \n".to_string();
        let call_init = generate_single_code(Call("Sys.init".to_string(), 0), "", 0);
        gencodes.push(set_sp + &call_init);
    }
    // If You Has not Overcome it RIGHT NOW?
    // When you can Overcome it?
    // You Win When You Overcome it RIGHT NOW
    // You Loss When You Loss Right Now.
    for (codes_, name) in codes.into_iter().zip(fnames) {
        for code in codes_ {
            let code = generate_single_code(code, name, idx);
            gencodes.push(code);
            idx += 1;
        }
    }
    return gencodes;
}

fn output_to_file(gencodes: Vec<String>, filename: &str) -> std::io::Result<()> {
    let mut file = File::create(filename)?;
    for code in gencodes.into_iter() {
        file.write(code.as_bytes())?;
        file.write("\r\n".as_bytes())?;
    }
    Ok(())
}


fn compile(path: &str, has_boot: bool) -> std::io::Result<()> {
    let path_obj = Path::new(path);
    // parser
    if path_obj.is_file() {
        let parser = Parser::new(path);
        let codes = parser.parse();
        let outpath = path.replace("vm", "asm");
        let fname = path_obj.file_stem().unwrap().to_str().unwrap();
        let vmcode = generate_codes(vec![codes], vec![fname], has_boot);
        output_to_file(vmcode, &outpath)?;
    } else if path_obj.is_dir() {
        let mut fnames = vec![];
        let mut codes = vec![];
        for entry in path_obj.read_dir()? {
            let entry = entry?;
            let filepath = entry.path();
            if filepath.extension().unwrap() == "vm" {
                let code = Parser::new(filepath.to_str().unwrap()).parse();
                let fname = filepath.file_stem().unwrap().to_str().unwrap().to_owned();
                codes.push(code);
                fnames.push(fname);
            }
        }
        let fnames_str: Vec<&str> = fnames.iter().map(|name| name.as_str()).collect();
        let vmcode = generate_codes(codes, fnames_str, has_boot);
        // generate code aside with directory
        let outpath = path_obj.with_extension("asm");
        output_to_file(vmcode, outpath.to_str().unwrap())?
    } else {
        println!("path `{}` is not vm file or directory contains vm files", path);
    }
    Ok(())
}

fn main() -> std::io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() == 2 {
        compile(&args[1], false)?;
    } else if args.len() == 3 {
        compile(&args[1], true)?;
    } else {
        println!("Usage: ./mvm /path/to/some_vm_file_or_directory [boot]");
    }
    Ok(())
}
